# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import collections
import ujson as json

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.base.book.util import pretty_book, human_readable_timestamp_record
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.kraken.book_builder import WsBookBuild as BookBuilder
from coin.exchange.kraken.kr_rest.product import KrakenProduct

ChannelIDMap = collections.namedtuple('ChannelIDMap', ['trade', 'book'])


class KrakenWsHandler(WebsocketHandler):
  def __init__(self, *, event_dist=None, feed_checker_config=None):
    super().__init__(event_dist=event_dist, feed_checker_config=feed_checker_config)
    self._channel_id_map = ChannelIDMap({}, {})

  def _debug_print(self, bids, asks):
    self._logger.info('\nbids: %s\n asks: %s',
                      json.dumps(bids, indent=2),
                      json.dumps(asks, indent=2))

  def on_depth(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())

    if isinstance(data, dict):
      # self._logger.info('book sub response: %s', data)
      channel_id = data.get('channelID')
      if channel_id is None:
        return

      pair = data['pair']
      product = KrakenProduct.FromStr(pair)
      self._channel_id_map.book[channel_id] = product
      return

    product = self._channel_id_map.book.get(data[0])
    if product is None:
      self._logger.error('Book data comes earlier than subscription response!')
      return

    native_symbol = product.native_symbol
    if not self.is_subscribed('book', native_symbol):
      return

    if 'as' in data[1] or 'bs' in data[1]:
      assert len(data) in (2, 4), data
      if product.native_symbol not in self._books:
        book = BookBuilder(native_symbol, book_builder_observer=self._book_builder_observer)
        self._books[product.native_symbol] = book
      else:
        book = self._books[product.native_symbol]
        book.clear_book()

      bids = data[1].get('bs', [])
      asks = data[1].get('as', [])
      book.snapshot(bids, asks, record.timestamp, record.topic_id)
      self.ready = True
    else:
      if self.ready is False:
        self._logger.error('Received update before snapshot!')
        return

      book = self._books[product.native_symbol]
      if len(data) in (2, 4):
        asks = data[1].get('a', [])
        bids = data[1].get('b', [])
      elif len(data) in (3, 5):
        asks = data[1]['a']
        bids = data[2]['b']
      else:
        raise ValueError('Invalid data with length %s\n%s', len(data), data)
      book.update(bids, asks, record.timestamp, record.topic_id)

    if self._feed_checker:
      self._feed_checker.on_book(product=book.product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())

    if isinstance(data, dict):
      channel_id = data.get('channelID')
      if channel_id is None:
        return

      pair = data['pair']
      product = KrakenProduct.FromStr(pair)
      self._channel_id_map.trade[channel_id] = product
      return

    product = self._channel_id_map.trade.get(data[0])
    if product is None:
      self._logger.error('Trade data comes earlier than subscription response!')
      return

    native_symbol = product.native_symbol
    if not self.is_subscribed('trade', native_symbol):
      return

    for trade_info in data[1]:
      price = float(trade_info[0])
      qty = float(trade_info[1])
      # Filter out those: fetched_time - exchange_time is too high. say, 10 secs is enough.
      if record.timestamp * 1e-9 - float(trade_info[2]) > 10:
        continue
      if trade_info[3] == 'b':
        side = TradeSide.TRADE_BUY_SIDE
      elif trade_info[3] == 's':
        side = TradeSide.TRADE_SELL_SIDE
      else:
        raise ValueError('Unknown side: %s', data)
      trade_event = create_market_trade_event(record.timestamp,
                                              native_symbol,
                                              price,
                                              qty,
                                              side,
                                              topic_id=record.topic_id)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)
      self.publish_event('trade', native_symbol, trade_event)
